
from show.show_traffic import *

from utils import *
import os
from show.method_set import *
from matplotlib.legend_handler import HandlerTuple
from pylab import *

smooth_window = 20


def draw_loss(loss_name):
    Fig = plt.figure(figsize=(20, 10))
    img_path = img_dir_path + '/' + loss_name + '.png'
    l_list = []
    m_list = []
    for method_id, method in enumerate(method_key_list):
        loss_path = os.path.join(root_path, method, METHOD_SET[method], 'state_dict.npy')
        print(loss_path)
        loss_data = np.load(loss_path, allow_pickle=True)[()]
        train_k = 'train_' + loss_name
        valid_k = 'valid_' + loss_name

        smoothed_loss_data = np.zeros_like(loss_data[train_k])
        for i in range(smoothed_loss_data.shape[0]):
            smoothed_loss_data[i] = np.mean(loss_data[train_k][max(i - smooth_window + 1, 0):i + 1])
        L1, = plt.plot(smoothed_loss_data, c='C' + str(method_id), label=method, linestyle="-")
        smoothed_loss_data = np.zeros_like(loss_data[valid_k])
        for i in range(smoothed_loss_data.shape[0]):
            smoothed_loss_data[i] = np.mean(loss_data[valid_k][max(i - smooth_window + 1, 0):i + 1])
        L2, = plt.plot(smoothed_loss_data, c='C' + str(method_id), label=method, linestyle="--")
        l_list.append((L1, L2))
        m_list.append(method)


    plt.ylabel(loss_name_dict[loss_name], fontweight=500, fontsize=40)
    plt.xlabel('No. of Iter.', fontweight=500, fontsize=40)
    plt.xticks(size=tick_size)
    plt.yticks(size=tick_size)
    # plt.legend(prop = {'size':20})
    # plt.legend()
    legend(l_list, m_list, handler_map={tuple: HandlerTuple(ndivide=None)},prop = {'size':20})
    Fig.savefig(img_path)
    plt.close()


img_dir_path =  dir_name
print(img_dir_path)
if not os.path.exists(img_dir_path):
    os.makedirs(img_dir_path)

if __name__ == '__main__':
    mpl.style.use('default')
    for loss in loss_list:
        draw_loss(loss)
